# ==== pipeline.py ====
from memory import Memory
from registers import RegisterFile
from alu import ALU
from control_unit import ControlUnit
from pipeline_registers import IF_ID, ID_EX, EX_MEM, MEM_WB
from forwarding_unit import ForwardingUnit
from hazard_detection import HazardDetection


class Pipeline:
    def __init__(self):
        self.memory = Memory()
        self.reg_file = RegisterFile()
        self.control = ControlUnit()
        self.alu = ALU()
        self.PC = 0
        self.if_id = IF_ID()
        self.id_ex = ID_EX()
        self.ex_mem = EX_MEM()
        self.mem_wb = MEM_WB()
        self.stall = False
        self.branch_taken = False
        self.flush = False
        self.cycle_count = 0
        self.instructions_executed = 0
        self.stalls = 0
        self.forwarding_used = 0

    def load_program(self, program):
        self.memory.load_program(program)

    def reset(self):
        self.PC = 0
        self.reg_file.reset()
        self.if_id = IF_ID()
        self.id_ex = ID_EX()
        self.ex_mem = EX_MEM()
        self.mem_wb = MEM_WB()
        self.stall = False
        self.branch_taken = False
        self.flush = False
        self.cycle_count = 0
        self.instructions_executed = 0
        self.stalls = 0
        self.forwarding_used = 0

    def IF_stage(self):
        if self.stall:
            return

        instruction = self.memory.read_inst(self.PC)
        self.if_id.instruction = instruction
        self.if_id.PC = self.PC
        self.if_id.valid = (instruction != 0)

        if not self.branch_taken:
            self.PC += 1

    def ID_stage(self):
        if self.stall:
            self.id_ex.valid = False
            self.id_ex.instr = 0  # 插入NOP
            return

        if not self.if_id.valid:
            self.id_ex.valid = False
            self.id_ex.instr = 0  # 插入NOP
            return

        instruction = self.if_id.instruction
        opcode = (instruction >> 26) & 0x3F

        self.control.decode(opcode)
        signals = self.control.signals

        rs = (instruction >> 21) & 0x1F
        rt = (instruction >> 16) & 0x1F
        rd = (instruction >> 11) & 0x1F
        funct = instruction & 0x3F
        imm = instruction & 0xFFFF
        imm = imm if imm < 0x8000 else imm - 0x10000

        read_data1 = self.reg_file.read(rs)
        read_data2 = self.reg_file.read(rt)

        # 检测加载使用冒险
        self.stall = HazardDetection.detect_data_hazard(self.id_ex, self.if_id)
        if self.stall:
            self.stalls += 1
            self.id_ex.valid = False
            self.id_ex.instr = 0
            return

        # 填充ID/EX流水线寄存器
        self.id_ex.control_signals = signals.copy()
        self.id_ex.PC = self.if_id.PC
        self.id_ex.read_data1 = read_data1
        self.id_ex.read_data2 = read_data2
        self.id_ex.imm = imm
        self.id_ex.rt = rt
        self.id_ex.rd = rd
        self.id_ex.rs = rs
        self.id_ex.funct = funct
        self.id_ex.write_reg = rd if signals["RegDst"] else rt  # 设置目标寄存器
        self.id_ex.instr = instruction  # 存储原始指令
        self.id_ex.valid = True

        self.if_id.valid = False

    def EX_stage(self):
        if not self.id_ex.valid:
            self.ex_mem.valid = False
            self.ex_mem.instr = 0  # NOP
            return

        signals = self.id_ex.control_signals

        # 前递解决数据冒险
        operand1, forward_type1 = ForwardingUnit.forward_rs(self.id_ex, self.ex_mem, self.mem_wb)
        operand2, forward_type2 = ForwardingUnit.forward_rt(self.id_ex, self.ex_mem, self.mem_wb)

        if forward_type1 > 0 or forward_type2 > 0:
            self.forwarding_used += 1

        alu_operand2 = operand2 if not signals["ALUSrc"] else self.id_ex.imm

        # 确定ALU操作
        alu_control = self._get_alu_control(signals["ALUOp"], self.id_ex.funct)

        # 执行ALU操作
        alu_result, zero = self.alu.execute(alu_control, operand1, alu_operand2)

        # 填充EX/MEM流水线寄存器
        self.ex_mem.control_signals = signals
        self.ex_mem.ALU_result = alu_result
        self.ex_mem.write_data = operand2
        self.ex_mem.write_reg = self.id_ex.write_reg
        self.ex_mem.zero = zero
        self.ex_mem.instr = self.id_ex.instr  # 传递指令
        self.ex_mem.valid = True

        self.id_ex.valid = False

    def _get_alu_control(self, alu_op, funct):
        if alu_op == "00":  # lw/sw/addi
            return "ADD"
        elif alu_op == "01":  # beq
            return "SUB"
        else:  # R-type
            if funct == 0x20:
                return "ADD"  # add
            elif funct == 0x22:
                return "SUB"  # sub
            elif funct == 0x24:
                return "AND"  # and
            elif funct == 0x25:
                return "OR"  # or
            elif funct == 0x2A:
                return "SLT"  # slt
            else:
                return "ADD"

    def MEM_stage(self):
        if not self.ex_mem.valid:
            self.mem_wb.valid = False
            self.mem_wb.instr = 0  # NOP
            return

        signals = self.ex_mem.control_signals

        # 处理分支冒险
        self.branch_taken = HazardDetection.detect_control_hazard(self.ex_mem)
        if self.branch_taken:
            branch_target = self.ex_mem.PC + 1 + self.id_ex.imm
            self.PC = branch_target
            self.flush = True

        mem_read_data = 0
        if signals["MemRead"]:
            addr = self.ex_mem.ALU_result
            mem_read_data = self.memory.read_data(addr)
        elif signals["MemWrite"]:
            addr = self.ex_mem.ALU_result
            self.memory.write_data(addr, self.ex_mem.write_data)

        # 填充MEM/WB流水线寄存器
        self.mem_wb.control_signals = signals
        self.mem_wb.ALU_result = self.ex_mem.ALU_result
        self.mem_wb.mem_read_data = mem_read_data
        self.mem_wb.write_reg = self.ex_mem.write_reg
        self.mem_wb.instr = self.ex_mem.instr  # 传递指令
        self.mem_wb.valid = True

        self.ex_mem.valid = False

    def WB_stage(self):
        if not self.mem_wb.valid:
            return

        signals = self.mem_wb.control_signals

        # 所有指令到达WB阶段都计数
        self.instructions_executed += 1

        if signals["RegWrite"]:
            if signals["MemtoReg"]:
                result = self.mem_wb.mem_read_data
            else:
                result = self.mem_wb.ALU_result
            self.reg_file.write(self.mem_wb.write_reg, result)

    def flush_pipeline(self):
        if self.flush:
            self.if_id.valid = False
            self.id_ex.valid = False
            self.flush = False

    def run_cycle(self):
        self.cycle_count += 1
        self.WB_stage()
        self.MEM_stage()
        self.EX_stage()
        self.ID_stage()
        self.IF_stage()
        self.flush_pipeline()

        if self.stall:
            self.stall = False
        else:
            self.branch_taken = False

    def print_state(self):
        # 获取各阶段指令
        if_instr = f"0x{self.if_id.instruction:08X}" if self.if_id.valid and self.if_id.instruction != 0 else "NOP"
        id_instr = f"0x{self.id_ex.instr:08X}" if self.id_ex.valid and self.id_ex.instr != 0 else "NOP"
        ex_instr = f"0x{self.ex_mem.instr:08X}" if self.ex_mem.valid and self.ex_mem.instr != 0 else "NOP"
        mem_instr = f"0x{self.mem_wb.instr:08X}" if self.mem_wb.valid and self.mem_wb.instr != 0 else "NOP"
        wb_state = "WB" if self.mem_wb.valid else "-"

        print(f"{self.cycle_count:<5}{if_instr:<15}{id_instr:<15}{ex_instr:<15}{mem_instr:<15}{wb_state:<15}")

    def run(self, max_cycles=20, verbose=True):
        if verbose:
            print("五级流水线仿真开始")
            print(f"{'周期':<5}{'IF':<10}{'ID':<10}{'EX':<10}{'MEM':<10}{'WB':<10}")

        for _ in range(max_cycles):
            self.run_cycle()

            if verbose:
                self.print_state()

            # 检查是否完成 (所有流水线寄存器为空且PC超出程序范围)
            if not any([self.if_id.valid, self.id_ex.valid,
                        self.ex_mem.valid, self.mem_wb.valid]) and self.PC >= len(self.memory.inst_mem):
                break

        if verbose:
            print("\n==== 仿真结果 ====")
            print(f"总时钟周期数: {self.cycle_count}")
            print(f"执行指令数: {self.instructions_executed}")
            print(f"停顿周期数: {self.stalls}")
            print(f"前递使用次数: {self.forwarding_used}")
            cpi = self.cycle_count / self.instructions_executed if self.instructions_executed > 0 else 0
            print(f"CPI: {cpi:.2f}")

            print("\n寄存器状态:")
            for i in range(0, 32, 4):
                regs = [f"${i + j}: {self.reg_file.registers[i + j]}" for j in range(4)]
                print("\t".join(regs))